import numpy as np
import pandas as pd
import rrcf
from sklearn.ensemble import IsolationForest
import matplotlib as mpl
mpl.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['font.size'] = 13
# the following two line for Chinese fonts output setting
# plt.rcParams['font.sans-serif']=['SimHei']
# plt.rcParams['axes.unicode_minus'] = False

# Read data
OilWell = pd.read_csv('tp152.csv', index_col=0)
xticks_date = OilWell.index
# print(xticks_date)
OilWell.index = pd.to_datetime(OilWell.index)
# print(type(OilWell.index.values))
# print(OilWell.index.values)

data = OilWell['pipe'].astype(float).values

# Create events
events = {
'event #1'          : ('2020-1-03 16:25:00',
                      '2020-1-03 16:30:00'),
'event #2'          : ('2020-1-04 12:28:00',
                      '2020-1-04 12:38:00'),
'event #3'          : ('2020-1-04 23:28:00',
                      '2020-1-04 23:33:00'),
'event #4'          : ('2020-1-05 05:24:00',
                      '2020-1-05 05:30:00'),
'event #5'          : ('2020-1-05 06:53:00',
                      '2020-1-05 06:58:00'),
'event #6'          : ('2020-1-05 09:28:00',
                      '2020-1-05 09:33:00'),
'event #7'          : ('2020-1-05 11:23:00',
                      '2020-1-05 11:28:00'),
'event #8'          : ('2020-1-06 08:55:00',
                      '2020-1-06 09:00:00'),
'event #9'          : ('2020-1-07 00:36:00',
                      '2020-1-07 00:41:00')
}
OilWell['event'] = np.zeros(len(OilWell))
for event, duration in events.items():
    start, end = duration
    OilWell.loc[start:end, 'event'] = 1

# Next, we run the RRCF algorithm and compute the CoDisp for each point.
print("Begin computing  CoDisp values.")
# Set tree parameters
num_trees = 200
shingle_size = 30
tree_size = 500

# Use the "shingle" generator to create rolling window
points = rrcf.shingle(data, size = shingle_size)
points = np.vstack([point for point in points])
n = points.shape[0]
sample_size_range = (n // tree_size, tree_size)

forest = []
while len(forest) < num_trees:
    ixs = np.random.choice(n, size = sample_size_range, replace = False)
    trees = [rrcf.RCTree(points[ix], index_labels=ix) for ix in ixs]
    forest.extend(trees)

avg_codisp = pd.Series(0.0, index=np.arange(n))
index = np.zeros(n)

for tree in forest:
    codisp = pd.Series({leaf: tree.codisp(leaf) for leaf in tree.leaves})
    avg_codisp[codisp.index] += codisp
    np.add.at(index, codisp.index.values, 1)

avg_codisp /= index
avg_codisp.index = OilWell.iloc[(shingle_size - 1):].index

# Running constructing Isolation Forest
contamination = OilWell['event'].sum()/len(OilWell)
IF = IsolationForest(n_estimators=num_trees, max_samples=500,
                     contamination=contamination, random_state = 0, behaviour='new')
IF.fit(points)
if_scores = IF.score_samples(points)

if_scores = pd.Series(-if_scores,
                      index=(OilWell
                             .iloc[(shingle_size - 1):]
                             .index))
# Plotting the results
# Normalize anomaly scores to (0, 1)
avg_codisp = ((avg_codisp - avg_codisp.min())
              / (avg_codisp.max() - avg_codisp.min()))
if_scores = ((if_scores - if_scores.min())
              / (if_scores.max() - if_scores.min()))
# print("if_scores: \n", if_scores.shape, '\n', if_scores)
# np.savetxt('output_result.txt', test_result)
np.savetxt('scores_result.txt ', np.hstack([avg_codisp, if_scores]))

#
fig, ax = plt.subplots(2, figsize=(15, 6))

(OilWell['pipe'] / 1.0).plot(ax=ax[0], color='0.2', alpha=0.9, label='Pipeline pressure')
if_scores.plot(ax=ax[1], color='#7EBDE6', alpha=0.95, label='Isolation forest')
avg_codisp.plot(ax=ax[1], color='#E8685D', alpha=0.95, label='Robust random cut forest')
ax[1].legend(loc='upper left', fontsize=10)

for event, duration in events.items():
    start, end = duration
    ax[0].axvspan(start, end, alpha=0.5, color='springgreen')

ax[0].set_xlabel('')
ax[0].set_ylabel('Pipeline pressure(Mpa)', size=12)
ax[1].set_xlabel('Datetime', size=12)
ax[1].set_ylabel('Normalized anomaly score', size=12)

ax[0].grid(alpha=0.3)
ax[0].set_title('Oil-well pipeline pressure', size=14)
ax[0].xaxis.set_ticklabels([])
ax[0].legend(loc='lower left', fontsize=12)
ax[0].set_xlim(OilWell.index[0], OilWell.index[-1])
ax[1].set_xlim(OilWell.index[0], OilWell.index[-1])
ax[1].xaxis.set_ticklabels(['2020-1-03', '2020-1-04', '2020-1-04',
                            '2020-1-05', '2020-1-05', '2020-1-06',
                            '2020-1-06', '2020-1-07'])

plt.grid(alpha=0.3)
plt.axhline(y = 0.48, linestyle='--', linewidth = 1.0, color='blue', alpha=0.9)
plt.axhline(y =0.33, linestyle='-.',linewidth = 1.0, color='limegreen', alpha=0.9)
plt.tight_layout()
plt.savefig("fig_4.png", dpi=300)
plt.savefig("fig_4.svg", dpi=300)
plt.savefig("fig_4.pdf", dpi=300)
plt.savefig("fig_4.tif", dpi=300)

plt.show()
